from datetime import datetime
import logging
import os
import signal
from multiprocessing import Pool
from functools import partial
from itertools import product
import json

from ortools.constraint_solver import pywrapcp

from .ortools_data_converter import DataConverter
from .heuristics_matching import first_solution_matching, local_search_matching
from .constants import OPT_PATH


class Optimizer:
    def __init__(self, config_path, opt_back_path=OPT_PATH):
        self.logger = logging.getLogger('Optimizer')
        self.logger.debug(f"Initializing Optimizer")
        try:
            with open(os.path.join(opt_back_path, config_path), 'r') as f:
                optimization_params = json.load(f)
                self.data_converter_config_path = optimization_params['data_converter_config_path']
                self.max_workday_length = optimization_params['max_workday_length']
                self.max_wait_time = optimization_params['max_wait_time']
                self.time_limit = optimization_params['time_limit']
                self.apply_time_windows = optimization_params['apply_time_windows']
                self.first_solution_strategies = optimization_params['first_solution_strategies']
                self.local_search_metaheuristics = optimization_params['local_search_metaheuristics']
                self.print_solution = optimization_params['print_solution']
                self.order_penalty = optimization_params['order_penalty']
        except Exception as e:
            self.logger.error(f"Couldn't load config {config_path}, {e}")
            raise FileNotFoundError("Couldn't load config")
        self.data_converter = DataConverter(config_path=self.data_converter_config_path, opt_back_path=opt_back_path)

    def timeout_handler(self, signum, frame):
        """Handler used to interrupt optimization if it hangs"""
        self.logger.error(f'Signal handler called with signal {signum}')
        raise TimeoutError

    def optimize(self, data, first_solution_strategy, local_search_metaheuristic, time_limit=60, capacity=14):
        """Solve the VRP problem for data with chosen heuristic and metaheuristic"""
        route_id = str(first_solution_strategy) + '-' + str(local_search_metaheuristic)
        manager = pywrapcp.RoutingIndexManager(len(data['time_matrix']), data['num_vehicles'],
                                               data['starts'], data['ends'])
        routing = pywrapcp.RoutingModel(manager)

        def time_callback(from_index, to_index):
            """Returns the travel time between two nodes"""
            # Convert from routing variable Index to time matrix NodeIndex
            from_node = manager.IndexToNode(from_index)
            to_node = manager.IndexToNode(to_index)
            return data['time_matrix'][from_node, to_node]

        # Add time dimension
        transit_callback_index = routing.RegisterTransitCallback(time_callback)
        routing.SetArcCostEvaluatorOfAllVehicles(transit_callback_index)
        time = 'Time'
        routing.AddDimension(
            evaluator_index=transit_callback_index,
            slack_max=self.max_wait_time,  # allow waiting time
            capacity=self.max_workday_length,  # maximum time per vehicle
            fix_start_cumul_to_zero=True,  # force start cumul to zero
            name=time)
        time_dimension = routing.GetDimensionOrDie(time)

        # Set maximum number of orders for 1 courier
        node_dimension = 'node_count'
        routing.AddConstantDimension(
            value=1,  # increment by one for each node
            capacity=capacity + 1,  # max value of nodes in 1 route
            fix_start_cumul_to_zero=True,  # set count to zero
            name=node_dimension)

        # Add time window constraints for each location except depots
        if self.apply_time_windows:
            self.logger.debug(f"[{route_id}] Adding time window constraints")
            for location_idx, time_window in enumerate(data['time_windows']):
                if location_idx in data['starts']:
                    continue
                index = manager.NodeToIndex(location_idx)
                time_dimension.CumulVar(index).SetRange(time_window[0], time_window[1])

        # Add time window constraints for each vehicle start node
        self.logger.debug(f"[{route_id}] Adding time window constraints for start nodes")
        for vehicle_id in range(data['num_vehicles']):
            index = routing.Start(vehicle_id)
            time_dimension.CumulVar(index).SetRange(data['time_windows'][vehicle_id][0],
                                                    data['time_windows'][vehicle_id][1])

        # Add dimension to minimize (total work time)
        self.logger.debug(f"[{route_id}] Adding function to minimize")
        for i in range(data['num_vehicles']):
            routing.AddVariableMinimizedByFinalizer(
                time_dimension.CumulVar(routing.Start(i)))
            routing.AddVariableMinimizedByFinalizer(
                time_dimension.CumulVar(routing.End(i)))

        # Initialize optimization parameters
        self.logger.debug(f"[{route_id}] Initializing search parameters")
        search_parameters = pywrapcp.DefaultRoutingSearchParameters()
        # First solution strategies/local metaheuristics
        # https://developers.google.com/optimization/routing/routing_options
        search_parameters.first_solution_strategy = first_solution_strategy
        search_parameters.local_search_metaheuristic = local_search_metaheuristic
        search_parameters.use_full_propagation = False  # Search for optimal solutions, affects speed tremendously
        search_parameters.time_limit.seconds = time_limit  # Time limit for early stopping
        # search_parameters.solution_limit = 10 # Found solutions count for early stopping
        search_parameters.log_search = False  # Log optimization process to console

        # Solve optimization task and convert results to human-readable route
        self.logger.debug(f"[{route_id}] Finding solution")
        assignment = routing.SolveWithParameters(search_parameters)
        if assignment:
            self.logger.debug(f"[{route_id}] Solution found, building route")
            if self.print_solution:
                self.data_converter.print_solution(data, manager, routing, assignment, route_id)
            route = self.data_converter.assignment_to_route(data, manager, routing, assignment, route_id)
            self.logger.debug(f"[{route_id}] Route built")
            return route
        else:
            return None

    def get_route(self, data, time_limit, capacity, heuristics):
        """Wrapper around optimize function"""
        first_solution_strategy, local_search_metaheuristic = heuristics
        route_id = str(first_solution_strategy) + '-' + str(local_search_metaheuristic)

        self.logger.info(
            f'[{route_id}] Trying to find the solution with {first_solution_matching[first_solution_strategy]} and '
            f'{local_search_matching[local_search_metaheuristic]} in {time_limit} seconds')
        self.logger.debug(f"[{route_id}] Started optimization")
        start = datetime.now()
        try:
            signal.signal(signal.SIGALRM, self.timeout_handler)
            signal.alarm(time_limit + 5)  # Terminate optimization if it takes more than time_limit+3 seconds
            route = self.optimize(data=data,
                                  first_solution_strategy=first_solution_strategy,
                                  local_search_metaheuristic=local_search_metaheuristic,
                                  time_limit=time_limit,  # Metaheuristic time limit in seconds
                                  capacity=capacity)
        except TimeoutError:
            self.logger.error(f"[{route_id}] Optimization exceeded time limit")
            route = {'id': route_id, 'is_valid': False}
            return route
        except Exception as e:
            self.logger.error(f"[{route_id}] Optimization failed, error: {e}")
            signal.alarm(0)
            route = {'id': route_id, 'is_valid': False}
            return route
        finally:
            signal.alarm(0)

        optimization_time = datetime.now() - start
        self.logger.debug(f"[{route_id}] Optimization took {optimization_time}")
        if route:
            self.logger.info(f"[{route_id}] Total route travel time: {route['total_travel_time']}min, "
                             f"total work time: {route['total_work_time']}min")
            route['id'] = route_id
            route['is_valid'] = True
        else:
            self.logger.warning(f"[{route_id}] No solution was found")
            route = {'id': route_id, 'is_valid': False}
        return route

    def get_routes(self, req_json, data=None, capacity=14):
        """Get a list of solutions using combinations of different heuristics and metaheuristics"""
        if not data:
            data, _ = self.data_converter.prepare_data(req_json)
        # Call optimization for each heuristic/metaheuristic combination
        partial_func = partial(self.get_route, data, self.time_limit, capacity)
        heuristics = list(product(self.first_solution_strategies, self.local_search_metaheuristics))
        self.logger.debug(f'Using {min(os.cpu_count(), len(heuristics))} parallel processes')
        with Pool(min(os.cpu_count(), len(heuristics))) as p:
            routes = p.map(partial_func, heuristics)
        return routes

    def pick_best_route(self, routes):
        best_route = None
        best_time = 99999
        max_node_count = 0
        for route in routes:
            if route['is_valid']:
                curr_node_count = sum([len(subroute['points']) for subroute in route['subroutes']])
                logging.info(f"[{route['id']}] {curr_node_count} nodes, "
                             f"{route['total_work_time']} total_work_time")
                if curr_node_count > max_node_count:
                    best_route = route
                    best_time = route['total_work_time']
                    max_node_count = curr_node_count
                elif route['total_work_time'] < best_time and curr_node_count == max_node_count:
                    best_route = route
                    best_time = route['total_work_time']
            else:
                if 'subroutes' in route:
                    self.logger.warning(f"[{route['id']}] Route is not valid")
                else:
                    self.logger.warning(f"[{route['id']}] Solution not found")
        if best_route:
            logging.info(f'Best route is {best_route["id"]}')
        else:
            logging.warning(f'No best route')
        return best_route
